# import numpy as np
import cv2
import mediapipe as mp

class PoseProcessor:
    """姿态处理类（基于MediaPipe）"""

    # def __init__(self):
    #     self.mp_pose = mp.solutions.pose  # 从MediaPipe里获取姿态估计的工具集
    #     # 创建一个姿态检测器对象（核心工具）
    #     self.pose = self.mp_pose.Pose(
    #         model_complexity=1,  # 模型复杂度（0-2，越高越准但越慢）
    #         min_detection_confidence=0.5,  # 检测置信度阈值：只有置信度>0.5才认为检测到了
    #         min_tracking_confidence=0.5  # 跟踪置信度阈值：跟踪时置信度>0.5才认为跟踪有效
    #     )
    #     self.mp_drawing = mp.solutions.drawing_utils  # MediaPipe的绘图工具，用来画关键点
    #
    # def process(self, frame):  # frame是输入的图像（比如摄像头拍的一帧画面）
    #     # 第一步：转换图像格式
    #     frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    #     # OpenCV读的图默认是BGR格式，而MediaPipe需要RGB格式，所以这里转换一下
    #
    #     # 第二步：用姿态检测器处理图像
    #     results = self.pose.process(frame_rgb)
    #     # 这里是核心！让前面初始化的pose模型去分析RGB图像，得到检测结果results
    #     # results里可能包含人体的关键姿态点坐标（如果检测到人的话）
    #
    #     # 第三步：如果检测到了姿态点，就画在原图上
    #     if results.pose_landmarks:  # 如果results里有姿态点（pose_landmarks不为空）
    #         self.mp_drawing.draw_landmarks(
    #             frame,  # 要画的图像（原图）
    #             results.pose_landmarks,  # 姿态点坐标数据
    #             self.mp_pose.POSE_CONNECTIONS,  # 哪些点之间需要连线（比如肩膀到手肘）
    #             # 关键点的样式：红色，线粗2，点的半径3
    #             self.mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2, circle_radius=3),
    #             # 连接线的样式：绿色，线粗2
    #             self.mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2)
    #         )
    #     # 第四步：返回结果
    #     return frame, results
    #     # frame：画好关键点的图像（可以直接显示或保存）
    #     # results：检测的原始结果（里面有所有关键点的坐标等详细信息）
    def __init__(self):
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            model_complexity=1,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        )
        self.mp_drawing = mp.solutions.drawing_utils

    def extract_landmark_coords(self, results, frame_shape):
        """
        从检测结果中提取关键点坐标（支持归一化坐标和像素坐标）
        参数：
            results: mediapipe的检测结果对象
            frame_shape: 图像的形状 (height, width)，用于将归一化坐标转换为像素坐标
        返回：
            coords: 字典，格式 {关键点ID: (x_norm, y_norm, z_norm, x_pixel, y_pixel)}
                - x_norm, y_norm, z_norm: 归一化坐标（0-1范围，相对于图像）
                - x_pixel, y_pixel: 像素坐标（实际图像中的像素位置）
        """
        coords = {}
        if results.pose_landmarks:
            frame_height, frame_width = frame_shape  # 获取图像的高和宽
            for idx, landmark in enumerate(results.pose_landmarks.landmark):
                # 归一化坐标（MediaPipe默认输出，x/y范围0-1，z为深度相对值）
                x_norm = landmark.x
                y_norm = landmark.y
                z_norm = landmark.z

                # 转换为像素坐标（x_pixel = x_norm * 图像宽度，y_pixel = y_norm * 图像高度）
                x_pixel = int(x_norm * frame_width)
                y_pixel = int(y_norm * frame_height)

                # 存入字典（关键点ID为key，包含5个值）
                coords[idx] = (x_norm, y_norm, z_norm, x_pixel, y_pixel)
        return coords

    def process(self, frame):
        """
        处理图像，返回绘制后的图像、原始检测结果、关键点坐标
        参数：
            frame: 输入图像（BGR格式，OpenCV默认格式）
        返回：
            frame: 绘制了关键点的图像
            results: mediapipe原始检测结果
            coords: 关键点坐标字典（同extract_landmark_coords的返回格式）
        """
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = self.pose.process(frame_rgb)

        # 绘制关键点（保持原有功能）
        if results.pose_landmarks:
            self.mp_drawing.draw_landmarks(
                frame,
                results.pose_landmarks,
                self.mp_pose.POSE_CONNECTIONS,
                self.mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2, circle_radius=3),
                self.mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2)
            )

        # 提取坐标（新增功能：传入图像形状，获取像素坐标）
        frame_shape = (frame.shape[0], frame.shape[1])  # (height, width)
        coords = self.extract_landmark_coords(results, frame_shape)

        # 同时返回绘制后的图像、原始结果、坐标字典
        return frame, results, coords
